import time
import selenium
from selenium import webdriver
from selenium.webdriver.common.by import By
import pandas as pd

# 该脚本的作用是读取已有的problem.xlsx，对给定范围内的题号，如果不在problem.xlsx中则进行爬取，最后将新爬取的写入另一个文件
# 该脚本开发原因是第一版脚本没有写异常捕获，导致没有标签/难度的题目爬取失败

# 可配置部分

exiting_problem_route = "combined_problems_infoB2.xlsx"
destin_problem_route = "problems_info_updatedB3.xlsx"
prefix = 'P'
left = 1000
right = 9999

# end
# 创建一个空的数据结构
data = {
    "Problem Number": [],
    "Problem Name": [],
    "Difficulty": [],
    "Tags": [],
    "Submission Count": [],
    "Accepted Count": [],
    "Acceptance Rate": []
}

def parse_number(num_str):
    """解析带有 M 和 K 的字符串，转换为浮点数"""
    num_str = num_str.strip()  # 去除前后空格
    if not num_str:
        return 0.0
    last_char = num_str[-1]  # 最后一个字符
    new_str = num_str[:-1]  # 去掉最后一个字符的串
    try:
        if last_char in ['K', 'k']:
            return float(new_str) * 1000.0
        elif last_char in ['M', 'm']:
            return float(new_str) * 1000000.0
        else:
            return float(num_str)
    except ValueError:
        return 0.0

def get_element_text(xpath):
    """尝试获取元素文本，若失败则返回空字符串"""
    try:
        element = browser.find_element(By.XPATH, xpath)
        return element.text.strip() if element.text else ""
    except Exception:
        return ""

# 读取已存储的 Excel 表
try:
    existing_df = pd.read_excel(exiting_problem_route)
    existing_problems = existing_df["Problem Number"].tolist()
except FileNotFoundError:
    existing_problems = []

# 启动 Chrome 浏览器
browser = webdriver.Chrome()

# 遍历从 P1000 到 P11079 的题目
for i in range(left, right):
    problem_number = f"{prefix}{i}"

    # 检查题号是否已存在
    if problem_number in existing_problems:
        print(f"{problem_number} already exists, skipping.")
        continue

    default_url = f"https://www.luogu.com.cn/problem/{problem_number}"
    browser.get(default_url)
    browser.implicitly_wait(5)

    try:
        # 尝试点击展开按钮，如果不存在，则跳过点击
        try:
            button = browser.find_element(By.XPATH, "//*[@id='app']/div[2]/main/div/section[1]/div[2]/div[2]/span")
            browser.execute_script("arguments[0].click();", button)
        except Exception:
            # 如果按钮不存在，说明没有标签，继续执行
            pass

        # 获取问题标题和难度
        problem_title = get_element_text("//*[@id='app']/div[2]/div[1]/div[2]/h1/span")
        problem_diff = get_element_text("//*[@id='app']/div[2]/main/div/section[1]/div[1]/div/div[2]/span[2]/a/span")

        # 将标题分解为题号和题名
        if problem_title:
            try:
                problem_number, problem_name = problem_title.split(" ", 1)
            except ValueError:
                problem_name = problem_title
        else:
            problem_number, problem_name = "", ""

        # 获取多个标签
        tag_list = []
        try:
            tags = browser.find_elements(By.XPATH, "//*[@id='app']/div[2]/main/div/section[1]/div[2]/div[1]/div/a")
            for tag in tags:
                href = tag.get_attribute("href")
                if href:
                    tag_number = href.split("=")[-1]
                    tag_list.append(tag_number)
        except Exception:
            # 如果获取标签失败，设置为空列表
            pass

        # 将标签列表转换为字符串
        tags_str = ", ".join(tag_list) if tag_list else ""

        # 获取提交次数和通过次数
        submission_count_str = get_element_text("//*[@id='app']/div[2]/div[1]/div[2]/div[2]/div[1]/div/div[1]/span[2]")
        accepted_count_str = get_element_text("//*[@id='app']/div[2]/div[1]/div[2]/div[2]/div[1]/div/div[2]/span[2]")

        # 将提交和通过次数存为字符串
        submission_count = submission_count_str if submission_count_str else ""
        accepted_count = accepted_count_str if accepted_count_str else ""

        # 计算通过率
        Acceptance_Rate = ""
        if submission_count and accepted_count:
            try:
                parsed_submission_count = parse_number(submission_count)
                parsed_accepted_count = parse_number(accepted_count)
                if parsed_submission_count > 0:
                    Acceptance_Rate = parsed_accepted_count / parsed_submission_count
                else:
                    Acceptance_Rate = 0.0
            except Exception:
                Acceptance_Rate = 0.0

        # 将信息存储到数据结构中
        # data["Problem Number"].append(problem_number if problem_number else "")
        # data["Problem Name"].append(problem_name if problem_name else "")
        # data["Difficulty"].append(problem_diff if problem_diff else "")
        # data["Tags"].append(tags_str)
        # data["Submission Count"].append(submission_count)
        # data["Accepted Count"].append(accepted_count)
        # data["Acceptance Rate"].append(Acceptance_Rate)
        # 在添加信息到数据结构之前检查题号是否已存在
        if problem_number not in data["Problem Number"]:
            data["Problem Number"].append(problem_number if problem_number else "")
            data["Problem Name"].append(problem_name if problem_name else "")
            data["Difficulty"].append(problem_diff if problem_diff else "")
            data["Tags"].append(tags_str)
            data["Submission Count"].append(submission_count)
            data["Accepted Count"].append(accepted_count)
            data["Acceptance Rate"].append(Acceptance_Rate)
        else:
            print(f"{problem_number} has already been added, skipping.")

    except Exception as e:
        print(f"Error fetching problem {problem_number}: {e}")

    time.sleep(0.5)  # 等待0.5秒，避免过快请求

# 创建 DataFrame
df = pd.DataFrame(data)

# 去除题号为空的行
df = df[df["Problem Number"] != ""]

# 将 DataFrame 写入 Excel 文件
df.to_excel(destin_problem_route, index=False)

browser.quit()

